#!/opt/python2.7/bin/python
##
#    Pass Extract
#    Password manager class for Python
#    
#    @author     Theodore Salmon, tass2001 -at- gmail -dot- com
#    @package    pass-extract
#    @version    $Id$
import base64, hashlib, xml.dom.minidom
from Crypto.Cipher import AES

class pass_extract(object):

    def __init__(self,username,app,password=None):
        '''
            Load all the things. Setup our AES object for decryption
            This will set two object variables - current_pass and username which you can access to return the login data
            @param str  username    The username
            @param str  app         The name of the application
            @param str  password    The password for the user/app to add into the XML. Note that his is only required to update/add an existing entry.
        '''
        self.username = username
        self.xmlFile = 'userLoginData.xml'
        self.xmlData = xml.dom.minidom.parse(self.xmlFile)
        iv = hashlib.sha1(app).hexdigest()[24:]
        key = hashlib.sha256(username).hexdigest()
        self.cipher = AES.new(key[0:len(key)/2], AES.MODE_CBC, iv)
        if password:
            self._edit_log(username,app,password)
        else:
            self._get_login(username,app)
            
    def _get_login(self,username,app):
        '''
            Get our login info
            @param string username  The username
            @param string app       The Application to lookup
        '''
        for user in self.xmlData.getElementsByTagName('User'):
            if user.getAttribute('App') == app and user.getAttribute('Username') == username:
                passHash = base64.b64decode(user.getAttribute('passwordHash'))
                self.current_pass = self.cipher.decrypt(passHash).rstrip(chr(0))

    def _edit_log(self,username,app,password):
        '''
            Add the login data into the XML or update the entry if it exists
            @param string username  The username
            @param string app       The Application name
            @param string password  The password
        '''
        pad = lambda inString: inString + (32 - len(inString) % 32) * chr(0)
        newElement = False
        passwordHash = base64.b64encode(self.cipher.encrypt(pad(password)))
        for userData in self.xmlData.getElementsByTagName('User'):
            if userData.getAttribute('App') == app and userData.getAttribute('Username') == self.username:
                newElement = userData
        if not newElement:
            newElement = self.xmlData.createElement('User')
            newElement.setAttribute('Username',self.username)
            newElement.setAttribute('App',app)
            newElement.setAttribute('passwordHash', passwordHash)
            self.xmlData.getElementsByTagName('Users').item(0).appendChild(newElement)
        else:
            newElement.setAttribute('passwordHash', passwordHash)
        self.current_pass = password
        self._save_xml()
        
    def _save_xml(self):
        self.xmlData.writexml(open(self.xmlFile,'w'))